# Databricks notebook source
# MAGIC %md
# MAGIC # Installations

# COMMAND ----------

# MAGIC %pip install -U mostlyai[local] 
# MAGIC %pip install ipywidgets
# MAGIC %restart_python

# COMMAND ----------

import time

# COMMAND ----------

time.sleep(600)

# COMMAND ----------

import pandas as pd
from mostlyai.sdk import MostlyAI

mostly = MostlyAI(local=True)

# COMMAND ----------

# MAGIC %md
# MAGIC # Get generator

# COMMAND ----------

[print(i) for i in mostly.generators.list()]

# COMMAND ----------

generator = mostly.generators.get("82340210-c89b-4db8-9bd8-e032ac6a2c6b")

# COMMAND ----------

# MAGIC %md
# MAGIC ## Save generator

# COMMAND ----------

generator.export_to_file('/Volumes/admin_govern_sta_des/tmp_mostlyai/generator/dSocial_31_03.zip')

# COMMAND ----------

df_personas = spark.table("admin_govern_sta_des.tmp_mostlyai.personas")
#df_contactes = spark.table("admin_govern_sta_des.tmp_mostlyai.contactes")
df_documents_personals = spark.table("admin_govern_sta_des.tmp_mostlyai.documents_personals")
df_persona_adreca = spark.table("admin_govern_sta_des.tmp_mostlyai.`persona-adreca`")
df_adreca = spark.table("admin_govern_sta_des.tmp_mostlyai.adreca")

# COMMAND ----------

# MAGIC %md
# MAGIC # Generation

# COMMAND ----------

# MAGIC %md
# MAGIC ## Configure data generation

# COMMAND ----------

config = {
    "tables": [
        {
            "name": "personas",
            "configuration": {
                "sample_size": 3378516
            }
        },
        {
            "name": "adreca",
            "configuration": {
                "sample_seed_data": df_adreca
            }
        }
    ]
}

# COMMAND ----------

# MAGIC %md
# MAGIC ## Generate data

# COMMAND ----------

df_samples = mostly.generate(generator, config=config)